
import mindspore as ms

from mindspore import nn
from mindspore.nn import WithLossCell

class CustomWithTrainCell(WithLossCell):
    def __init__(self, backbone, loss_fn):
        super(CustomWithTrainCell, self).__init__(backbone, loss_fn)

    def construct(self, imgA, imgB, label):

        outputs = self._backbone(imgA, imgB)
        # print('execute to cal loss ...')
        # pred = outputs[-1]
        return self._loss_fn(outputs, label)



class CustomWithEvalCell(nn.Cell):

    def __init__(self, network, loss_fn):
        super(CustomWithEvalCell, self).__init__(auto_prefix=False)
        self.network = network
        self.loss_fn = loss_fn

    def construct(self, input1, input2, label):
        outputs = self.network(input1, input2)
        pred = outputs[-1]
        loss = self.loss_fn(outputs, label)
        return loss, pred
